# 在app03.py的基础上，增加对聊天记录的保存，并将历史记录的链接显示在左侧，点击后显示在聊天区。
# 个别代码修改： 添加chdir当前目录。将历史记录放在设置之下。去掉设置、历史对话等内容。
# 在app04改进：1.新建会话后，没有清空聊天区。2.load_history应该把目录下所有的记录都显示和链接。3.历史会话列表可以单独进行删减

"""
新建会话优化：新增st.rerun()强制刷新界面,明确重置messages和current_chat_id
历史记录管理增强：添加删除按钮和delete_chat_history函数,使用列布局实现记录+删除按钮的组合,删除后自动刷新界面
状态管理改进：新增force_refresh状态控制界面刷新, 使用st.rerun()确保状态同步, 改进时间显示格式
"""

# -*- coding: utf-8 -*-
import streamlit as st
import requests
import json
import os
import uuid
from datetime import datetime
import shutil

st.set_page_config(
    page_title="Qwen2.5 Chat",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

custom_css = """
<style>
[data-testid="stAppViewContainer"] { background-color: #f5f5f5; }
[data-testid="stHeader"] { background-color: #1a73e8; }
[data-testid="stSidebar"] { background-color: #ffffff; }
.chat-item { display: flex; justify-content: space-between; align-items: center; padding: 8px; }
.chat-link { cursor: pointer; flex-grow: 1; }
.delete-btn { color: #ff4b4b; cursor: pointer; margin-left: 10px; }
</style>
"""
st.markdown(custom_css, unsafe_allow_html=True)

SAVE_DIR = "saved_chats"
os.chdir(os.path.dirname(os.path.realpath(__file__)))   #新增代码
os.makedirs(SAVE_DIR, exist_ok=True)

SERVER_CONFIGS = [
    {
        "name": "Ollama",
        "url": "http://localhost:11434/v1",
        "models": ["qwen2.5:3b", "deepseek-r1:1.5b","phi", "glm4"],
    },
    {
        "name": "AiGitee",
        "url": "https://ai.gitee.com/v1",
        "models": ["Qwen2.5-7B-Instruct"]
    },
    {
        "name": "ThuDS",
        "url": "https://madmodel.cs.tsinghua.edu.cn/v1",
        "models": ["DeepSeek-R1-Distill-32B", "DeepSeek-R1-671B"],
    }
]

# 初始化session状态
if "settings" not in st.session_state:
    st.session_state.update({
        "settings": {
            "show_form": True,
            "selected_server": SERVER_CONFIGS[0]["name"],
            "api_url": SERVER_CONFIGS[0]["url"],
            "selected_model": SERVER_CONFIGS[0]["models"][0],
            "api_key": "",
            "current_chat_id": None
        },
        "messages": [],
        "force_refresh": False
    })

def save_chat_history():
    if st.session_state.messages:
        chat_id = st.session_state.settings["current_chat_id"] or str(uuid.uuid4())
        filename = f"{SAVE_DIR}/{chat_id}.json"
        
        chat_data = {
            "id": chat_id,
            "title": st.session_state.messages[0]["content"][:20] + "...",
            "timestamp": datetime.now().isoformat(),
            "messages": st.session_state.messages
        }
        
        with open(filename, "w", encoding="utf-8") as f:
            json.dump(chat_data, f, ensure_ascii=False)
        
        st.session_state.settings["current_chat_id"] = chat_id
        st.session_state.force_refresh = True

def delete_chat_history(chat_id):
    try:
        os.remove(f"{SAVE_DIR}/{chat_id}.json")
        st.session_state.force_refresh = True
        if st.session_state.settings["current_chat_id"] == chat_id:
            st.session_state.messages = []
            st.session_state.settings["current_chat_id"] = None
    except Exception as e:
        st.error(f"删除失败: {str(e)}")

def load_chat_history(chat_id):
    try:
        with open(f"{SAVE_DIR}/{chat_id}.json", "r", encoding="utf-8") as f:
            chat_data = json.load(f)
            st.session_state.messages = chat_data["messages"]
            st.session_state.settings["current_chat_id"] = chat_id
    except Exception as e:
        st.error(f"加载失败: {str(e)}")

def get_saved_chats():
    chats = []
    for fname in os.listdir(SAVE_DIR):
        if fname.endswith(".json"):
            try:
                with open(f"{SAVE_DIR}/{fname}", "r", encoding="utf-8") as f:
                    chat_data = json.load(f)
                    chats.append(chat_data)
            except:
                continue
    return sorted(chats, key=lambda x: x["timestamp"], reverse=True)

with st.sidebar:
    st.header("对话管理")

    # 配置表单（保持原有代码不变）...
    if st.button("显示/隐藏配置"):
        st.session_state.settings["show_form"] = not st.session_state.settings["show_form"]
    
    if st.session_state.settings["show_form"]:
        with st.form("config_form"):
            new_server = st.selectbox(
                "服务地址",
                options=[c["name"] for c in SERVER_CONFIGS],
                index=next(i for i, c in enumerate(SERVER_CONFIGS) if c["name"] == st.session_state.settings["selected_server"])
            )
            
            if new_server != st.session_state.settings["selected_server"]:
                st.session_state.settings["selected_server"] = new_server
                server_config = next(c for c in SERVER_CONFIGS if c["name"] == new_server)
                st.session_state.settings["api_url"] = server_config["url"]
                st.session_state.settings["selected_model"] = server_config["models"][0]
            
            current_models = next(c["models"] for c in SERVER_CONFIGS if c["name"] == st.session_state.settings["selected_server"])
            st.session_state.settings["selected_model"] = st.selectbox(
                "模型选择",
                options=current_models,
                index=current_models.index(st.session_state.settings["selected_model"])
            )
            
            st.session_state.settings["api_key"] = st.text_input(
                "API密钥",
                type="password",
                value=st.session_state.settings["api_key"]
            )
            
            st.form_submit_button("保存配置")
    
    # 新建会话
    if st.button("新建会话"):
        save_chat_history()
        st.session_state.messages = []
        st.session_state.settings["current_chat_id"] = None
        st.rerun()
    
    # 历史记录列表
    st.subheader("历史会话")
    saved_chats = get_saved_chats()
    
    for chat in saved_chats:
        col1, col2 = st.columns([4,1])
        with col1:
            if st.button(
                f"{chat['title']}\n{datetime.fromisoformat(chat['timestamp']).strftime('%Y-%m-%d %H:%M')}",
                key=f"btn_{chat['id']}",
                use_container_width=True
            ):
                load_chat_history(chat["id"])
        
        with col2:
            if st.button("×", key=f"del_{chat['id']}"):
                delete_chat_history(chat["id"])
                st.rerun()


# 消息显示和生成逻辑（保持原有代码不变）...  
def generate_response(prompt):
    try:
        headers = {
            "Authorization": f"Bearer {st.session_state.settings['api_key']}",
            "Content-Type": "application/json"
        } if st.session_state.settings["api_key"] else {}
        
        response = requests.post(
            f"{st.session_state.settings['api_url']}/chat/completions",
            headers=headers,
            json={
                "model": st.session_state.settings["selected_model"],
                "messages": [{'role':'user', 'content': prompt}],
                "stream": False
            },
            timeout=60
        )
        return json.loads(response.text)["choices"][0]["message"]["content"]
    except Exception as e:
        return f"请求出错: {str(e)}"

# 强制刷新处理
if st.session_state.force_refresh:
    st.session_state.force_refresh = False
    st.rerun()

# 加载URL参数指定的聊天记录
# query_params = st.experimental_get_query_params()  # 弃用
query_params = st.query_params
if "chat_id" in query_params:
    load_chat_history(query_params["chat_id"][0])

# 消息显示区域
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(f'<div class="{"user-message" if message["role"] == "user" else "bot-message"}">{message["content"]}</div>', 
                    unsafe_allow_html=True)

# 消息输入处理
if prompt := st.chat_input("输入你的消息..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    save_chat_history()
    
    with st.chat_message("user"):
        st.markdown(f'<div class="user-message">{prompt}</div>', unsafe_allow_html=True)

    with st.chat_message("assistant"):
        with st.spinner("正在思考..."):
            full_response = generate_response(prompt)
            st.markdown(f'<div class="bot-message">{full_response}</div>', unsafe_allow_html=True)
    
    st.session_state.messages.append({"role": "assistant", "content": full_response})
    save_chat_history()
    st.rerun()